#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2018/12/25 11:48
# @USER    : Connor
# @File    : submit_macLV.py
# @Software: PyCharm
# @Version  : Python-3.6
# @TASK:
import logging
import sys, os, cv2,glob,re
import numpy as np
import scipy.io as sio
from fcn_model import fcn_model
from helpers import center_crop, reshape

MAC_ROOT_PATH = 'MAC-Cine'

VALIDATION_PATH = os.path.join(MAC_ROOT_PATH)#, 'Validation')


def create_submission(dcm_list, data_path):
    logger.info('Loading FCN model...')
    if contour_type == 'i':
        weights = 'model_logs/sunnybrook_i_epoch_478_0.9905.h5'
    elif contour_type == 'o':
        weights = 'model_logs/sunnybrook_o_epoch_499_0.9862.h5'
    else:
        logger.error('\ncontour type "%s" not recognized\n' % contour_type)
        sys.exit('\ncontour type "%s" not recognized\n' % contour_type)

    crop_size = 100
    input_shape = (crop_size, crop_size, 1)
    num_classes = 2
    model = fcn_model(input_shape, num_classes, weights=weights)


    logger.info('Building Data Set...')
    for idx,dcm in enumerate(dcm_list):
        temp = load_CineImages(dcm)
        if idx==0:
            timages=temp
        else:
            timages = np.concatenate((timages,temp),axis=2)

    h,w,depth = timages.shape
    # h,w = timages.shape
    # depth=1
    images = np.zeros((depth, crop_size, crop_size, 1))

    # if timages.ndim < 3:
    #     img = timages[..., np.newaxis]
    # img = center_crop(img, crop_size=crop_size)
    # images[0] = img
    for idx in range(depth):
        img = timages[:,:,idx]
        if img.ndim < 3:
            img = img[..., np.newaxis]
        img = center_crop(img, crop_size=crop_size)
        images[idx] = img

    logger.info('Predicting...')
    pred_masks = model.predict(images, batch_size=32, verbose=1)

    save_dir = os.path.join(data_path,('FCN_auto_contour'))
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    np.save(os.path.join(save_dir,(contour_type+'_pred_masks.npy')),pred_masks)
    # output all results
    for idx,subject in enumerate(dcm_list):



    # prefix = 'MYFCN_'  # change prefix to your unique initials
    # for idx in enumerate(dcm_list):
    #     img = read_dicom(dcm_path)
    #     h, w, d = img.shape
    #     otmp = reshape(opred_masks[idx], to_shape=(h, w, d))
    #     otmp = np.where(otmp > 0.5, 255, 0).astype('uint8')
    #     itmp = reshape(ipred_masks[idx], to_shape=(h, w, d))
    #     itmp = np.where(itmp > 0.5, 255, 0).astype('uint8')
    #     assert img.shape == otmp.shape, 'Prediction does not match shape'
    #     assert img.shape == itmp.shape, 'Prediction does not match shape'
    #     tmp = otmp - itmp
    #     tmp = np.squeeze(tmp, axis=(2,))
    #     sub_dir = dcm_path[dcm_path.find('CAP_'):dcm_path.rfind('DET')]
    #     filename = prefix + dcm_path[dcm_path.rfind('DET'):].replace('.dcm', '.png')
    #     full_path = os.path.join(save_dir, sub_dir)
    #     if not os.path.exists(full_path):
    #         os.makedirs(full_path)
    #     cv2.imwrite(os.path.join(full_path, filename), tmp)
    #     in_ = cv2.imread(os.path.join(full_path, filename), cv2.IMREAD_GRAYSCALE)
    #     if not np.allclose(in_, tmp):
    #         raise AssertionError('File read error: {:s}'.format(os.path.join(full_path, filename)))


def load_CineImages(patientdir):
    # data = sio.loadmat(patientdir)
    # images = data['tt1']
    glob_search = os.path.join(patientdir, '*_sBTFE_BH_7_*.mat')
    matfiles = glob.glob(glob_search)
    if len(matfiles) == 0:
        logger.error("Couldn't find Cine-Short MAT file in {}. "
                     "Wrong directory?".format(patientdir))
        raise Exception
    for file in matfiles:
        data = sio.loadmat(file)
        # images = data['tt1']
        images = data['I']
        images = images[0, 0]
        images = np.array(images[0])

    return images


def get_all_dicoms(data_path):
    dcm_list=[os.path.join(data_path,p) for p in os.listdir(data_path) if 'P' in p]
    # dcm_list = [os.path.join(dirpath, f)
    #             for dirpath, dirnames, files in os.walk(data_path)
    #             for f in files if 'S' in f]
    logger.info('Number of examples: {:d}'.format(len(dcm_list)))
    # logger.info(dcm_list)

    return dcm_list


if __name__ == '__main__':
    logger = logging.getLogger()
    logger.setLevel(logging.NOTSET)  # Log等级总开关
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)
    formatter = logging.Formatter("%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s")
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    if len(sys.argv) < 2:
        logger.error('Usage: python %s <i/o> <gpu_id>' % sys.argv[0])
        sys.exit('Usage: python %s <i/o> <gpu_id>' % sys.argv[0])
    elif len(sys.argv) < 3:
        contour_type = sys.argv[1]
    else:
        contour_type = sys.argv[1]
        os.environ['CUDA_VISIBLE_DEVICES'] = sys.argv[2]

    logger.info('Processing Validation set ...')
    val_dicoms = get_all_dicoms(VALIDATION_PATH)
    print(val_dicoms)
    create_submission(val_dicoms, VALIDATION_PATH)
    logger.info('All done.')
